import pandas as pd
import csv
from agent import code_agent, cot, debate, advisor1, advisor2
from tool import *
import re
import json

COUNT = 0
ID = 0
TH = 4
DATA_TEMPLATE = {'id': 1, 'query': 'What is Python?', 'title': 'Python Basics', 'sample_label': 'Programming',
                 "label": 1, 'esci_label': ''}
Historical_fields = []


def extract_cot(text):
    # 使用正则表达式提取符合格式的整数和小数部分
    match = re.search(r'\[(\d+), (\d+\.\d+)\]', text)

    if match:
        integer_part = int(match.group(1))

        return integer_part

    else:
        return "ERROR"


def extract_code(text):
    try:

        score = float(json.loads(text)["scores"][0])
        if score > 0.5:
            return 1
        else:
            return 0
    except:
        print(f"extract_code error :{text}")
        return "ERROR"


def extract_advisor(text):
    try:
        result = eval(text.split('\'')[1])
        return result
    except:
        print(f"extract_advisor error :{text}")
        return "ERROR"


def extract_debate(text):
    try:
        matches = re.findall(r'(\d+)', text)
        last_match = matches[-1]
        return int(last_match)
    except:
        print(f"extract_debate error :{text}")
        return "ERROR"


def append_to_csv(filename, data):
    # 打开文件以追加模式写入
    with open(filename, 'a', newline='', encoding='utf-8') as csvfile:
        writer = csv.DictWriter(csvfile, fieldnames=data.keys())

        # 如果文件为空，写入列名
        if csvfile.tell() == 0:
            writer.writeheader()

        # 写入数据
        writer.writerow(data)


def process(field1, label, field2):
    global COUNT, ID, Historical_fields

    result_code = code_agent(field1, field2)
    result_cot = cot(field1, field2)



    e_code = extract_code(result_code)
    e_cot = extract_cot(result_cot)

    if e_cot == 1 and e_code == 1:
        # 积极样本
        data = DATA_TEMPLATE.copy()
        data['id'] = ID
        data['query'] = field1
        data['title'] = field2
        data['sample_label'] = "easy"
        data['label'] = e_cot
        data['esci_label'] = label
        append_to_csv("output.csv", data)
        COUNT += 1
        if COUNT == TH:
            Historical_fields = []
            ID += 1
            COUNT = 0
            return 0
        # look for hard sample
        new_query = advisor1(field2, Historical_fields).chat_history[1]["content"]
        Historical_fields.append(new_query)
        field1 = new_query
        process(field1, label, field2)
    elif e_cot == 0 and e_code == 0:
        # 消极样本
        data = DATA_TEMPLATE.copy()
        data['id'] = ID
        data['query'] = field1
        data['title'] = field2
        data['sample_label'] = "easy"
        data['label'] = e_cot
        data['esci_label'] = label
        append_to_csv("output.csv", data)
        COUNT += 1
        if COUNT == TH:
            Historical_fields = []
            ID += 1
            COUNT = 0
            return 0
        # look for hard sample
        new_query = advisor2(field2, Historical_fields).chat_history[1]["content"]
        Historical_fields.append(new_query)
        field1 = new_query
        process(field1, label, field2)
    else:
        data = DATA_TEMPLATE.copy()
        data['id'] = ID
        data['query'] = field1
        data['title'] = field2
        data['sample_label'] = "hard"

        debate_result = debate(field1, field2)
        e_debate = extract_debate(debate_result)

        data['label'] = e_cot if e_debate == "ERROR" else e_debate
        data['esci_label'] = label
        append_to_csv("output.csv", data)
        ID += 1


def main():
    df = pd.read_csv('output_utf8.csv')

    if 'query' not in df.columns or 'esci_label' not in df.columns or 'product_title' not in df.columns:
        print("CSV文件缺少必要的列名。请确保文件包含 'query', 'esci_label', 'product_title' 列。")
        return

    for index, row in df.iterrows():
        print("#" * 100)
        for i in range(5):
            print()
            if i == 2:
                print("query", row['query'])
                print("esci_label", row['esci_label'])
                print("product_title", row['product_title'])
        print("#" * 100)
        field1 = row['query']
        label = row['esci_label']
        field2 = row['product_title']

        process(field1, label, field2)


if __name__ == '__main__':
    main()
